#include <odp/odp.h>
#include <opencv2/opencv.hpp>

#include <iostream>

namespace odp {

    void ODP::SetParameter(float score_threshold, float iou_threshold) {
        this->score_threshold = score_threshold;
        this->iou_threshold = iou_threshold;
    }

    bool ODP::Process_Array(float *data_ptr, size_t data_len, int class_number) {
        int clos = (class_number + 5);
        if (data_len % clos != 0) return false;
        int rows = data_len / clos;
        cv::Mat data_mat = cv::Mat(rows, clos, CV_32FC1);
        for (int i = 0; i < rows; i++) {
            for (int j = 0; j < clos; j++) {
                data_mat.at<float>(i, j) = *(data_ptr + i * clos + j);
            }
        }
        return Process_Mat(data_mat);
    }

    bool ODP::Process_Mat(cv::Mat &data_mat) {
        cv::Mat score_result(0, data_mat.size().width, CV_32F);
        for (int i = 0; i < data_mat.size().height; ++i) {
            if (data_mat.at<float>(i, 4) > this->score_threshold) {
                score_result.push_back(data_mat.row(i));
            }
        }

        //  (中心 x, 中心 y, 宽度, 高度) 属性转换成 (左上角 x, 左上角 y, 右下角 x, 右下角 y)
        for (int i = 0; i < score_result.size().height; ++i) {
            float c_x = score_result.at<float>(i, 0);
            float c_y = score_result.at<float>(i, 1);
            score_result.at<float>(i, 0) = c_x - score_result.at<float>(i, 2) / 2.0;
            score_result.at<float>(i, 1) = c_y - score_result.at<float>(i, 3) / 2.0;
            score_result.at<float>(i, 2) = c_x + score_result.at<float>(i, 2) / 2.0;
            score_result.at<float>(i, 3) = c_y + score_result.at<float>(i, 3) / 2.0;
        }

        // get class 最大值索引, 最大值, 按照dim=1 方向计算
        cv::Mat class_result(0, 5 + 2, CV_32F);
        //只保留bbox4个坐标以及objectnness分数，转而增加了有最大值的类别分数及索引。
        for (int i = 0; i < score_result.size().height; ++i) {
            double minVal;
            double maxVal;
            int minIdx[2];
            int maxIdx[2];
            cv::minMaxIdx(score_result.row(i).colRange(5, score_result.size().width), &minVal, &maxVal, minIdx, maxIdx);
            cv::Mat tmp = (cv::Mat_<float>(1, 7) << score_result.at<float>(i, 0),
                    score_result.at<float>(i, 1),
                    score_result.at<float>(i, 2),
                    score_result.at<float>(i, 3),
                    score_result.at<float>(i, 4), // score
                    maxIdx[1], maxVal);
            class_result.push_back(tmp);
        }

        result_mat = cv::Mat(0, 5 + 2, CV_32F);
        // get all can used classes
        std::set<int> class_set = filters_set(class_result.col(5));
        for (auto &it: class_set) {
            cv::Mat box_result(0, 5 + 2, CV_32F);
            cv::Mat box_input(0, 5 + 2, CV_32F);
            for (int i = 0; i < class_result.size().height; ++i) {
                if ((int) class_result.at<float>(i, 5) == it) {
                    box_input.push_back(class_result.row(i));
                }
            }

            nms_functions(box_input, box_result);
            result_mat.push_back(box_result);
        }

        return true;
    }

    cv::Mat &ODP::GetResult() {
        return result_mat;
    }

    std::set<int> ODP::filters_set(cv::Mat m) {
        std::set<int> result;

        for (int i = 0; i < m.size().height; ++i) {
            result.insert((int) m.at<float>(i, 0));
        }
        return result;
    }

    void ODP::nms_functions(cv::Mat &input_mat, cv::Mat &result_mat) {
        cv::Mat boxs = input_mat.clone();

        while (boxs.size().height > 0) {
            double minVal;
            double maxVal;
            int minIdx[2];
            int maxIdx[2];
            cv::minMaxIdx(boxs.col(6), &minVal, &maxVal, minIdx, maxIdx);

            cv::Mat bast_box = boxs.row(maxIdx[0]);
            result_mat.push_back(bast_box);


            cv::Mat new_boxs(0, 5 + 2, CV_32F);
            for (int i = 0; i < boxs.size().height; ++i) {
                if (i == maxIdx[0]) continue;
                float iou = box_iou(bast_box, boxs.row(i));

                if (iou < this->iou_threshold) {
                    new_boxs.push_back(boxs.row(i));
                }
            }
            boxs = new_boxs;
        }
    }

    float ODP::box_iou(cv::Mat box1, cv::Mat box2) {

        float in_w = std::min(box1.at<float>(0, 2), box2.at<float>(0, 2)) -
                     std::max(box1.at<float>(0, 0), box2.at<float>(0, 0));
        float in_h = std::min(box1.at<float>(0, 3), box2.at<float>(0, 3)) -
                     std::max(box1.at<float>(0, 1), box2.at<float>(0, 1));

        float inter_area = in_w * in_h;

        float union_area = (box1.at<float>(0, 2) - box1.at<float>(0, 0)) *
                           (box1.at<float>(0, 3) - box1.at<float>(0, 1)) +
                           (box2.at<float>(0, 2) - box2.at<float>(0, 0)) *
                           (box2.at<float>(0, 3) - box2.at<float>(0, 1)) - inter_area;

        return inter_area / union_area;
    }


}


